
from typing import BinaryIO, List
import numpy as np
from PIL import Image
import torch

from monai.transforms import (
    EnsureChannelFirst,
    Compose,
    Transform,
    ScaleIntensity,
    EnsureType,
)

import bentoml
from bentoml.frameworks.pytorch import PytorchModelArtifact
from bentoml.adapters import FileInput, JsonOutput
from bentoml.utils import cached_property

MEDNIST_CLASSES = ["AbdomenCT", "BreastMRI", "CXR", "ChestCT", "Hand", "HeadCT"]


class LoadStreamPIL(Transform):
    """Load an image file from a data stream using PIL."""

    def __init__(self, mode=None):
        self.mode = mode

    def __call__(self, stream):
        img = Image.open(stream)

        if self.mode is not None:
            img = img.convert(mode=self.mode)

        return np.array(img)


@bentoml.env(pip_packages=["torch", "numpy", "monai", "pillow"])
@bentoml.artifacts([PytorchModelArtifact("classifier")])
class MedNISTClassifier(bentoml.BentoService):
    @cached_property
    def transform(self):
        return Compose([LoadStreamPIL("L"), EnsureChannelFirst(channel_dim="no_channel"), ScaleIntensity(), EnsureType()])

    @bentoml.api(input=FileInput(), output=JsonOutput(), batch=True)
    def predict(self, file_streams: List[BinaryIO]) -> List[str]:
        img_tensors = list(map(self.transform, file_streams))
        batch = torch.stack(img_tensors).float()

        with torch.no_grad():
            outputs = self.artifacts.classifier(batch)
        _, output_classes = outputs.max(dim=1)

        return [MEDNIST_CLASSES[oc] for oc in output_classes]
